#! /usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Author :   Liang Kang
@Contact:   gangkanli1219@163.com
@Time   :   2018/8/7 16:44
@Desc   :
"""
from abc import abstractmethod
from pathlib import Path

import numpy as np
import cv2

from dltools.data.dataset.cropper import BaseCropper
from dltools.utils.io import read_voc_xml
from dltools.utils.log import get_console_logger


class ShapeCropper(BaseCropper):

    def __init__(self, image, xml, output, image_file, exp_ratio=0.1,
                 logger=None, break_file=None):
        assert isinstance(image, str)
        if logger is not None:
            logger = logger.getChild('ShapeCropper')
        else:
            logger = get_console_logger('ShapeCropper')
        self._image_path = image
        self._xml_path = xml
        self._break_file = break_file
        self._output = output
        self._do_write = True
        super(ShapeCropper, self).__init__(image, logger)
        self._exp_ratio = exp_ratio
        self._image_file = image_file

    @abstractmethod
    def _preprocess(self):
        """
        预判断, 图像或标注文件是否符合要求

        Returns
        -------

        """
        self._image_info, break_instance = read_voc_xml(
            self._xml_path, self.image.shape)
        if break_instance or (len(self._image_info['objects']) == 0):
            self._logger.info('File {} is Broken !'.format(self._xml_path))
            if self._break_file is not None:
                with open(self._break_file, 'a', encoding='utf-8') as out:
                    out.write(self._image_path + '\n')
            return True
        return False

    @abstractmethod
    def _set_up(self):
        """
        确认输出路径
        Returns
        -------

        """
        if self._output is not None:
            self._output = Path(self._output)
            if not self._output.exists():
                self._output.mkdir(parents=True)
        else:
            self._do_write = False

    @abstractmethod
    def _crop_image(self):
        """
        切割出对象
        Returns
        -------

        """
        shape = self._buf_data['shape']
        boxes = self._buf_data['boxes']
        labels = self._buf_data['labels']
        min_axis, max_axis = np.split(boxes, 2, axis=1)
        size = (max_axis - min_axis) * self._exp_ratio
        size = np.ceil(size).astype(np.int32)
        min_axis = np.clip(min_axis - size, 0, shape)
        max_axis = np.clip(max_axis + size, 0, shape)
        boxes = np.hstack([min_axis, max_axis])

        for ob, lb in zip(boxes, labels):
            self._buf_data['box'] = ob
            self._buf_data['label'] = lb
            self._write_image()

    @abstractmethod
    def _write_image(self):
        """
        写入图像
        Returns
        -------

        """
        if not self._do_write:
            return
        box = self._buf_data['box']
        label = self._buf_data['label']
        sub_image = self.image[box[0]:box[2], box[1]:box[3], :]
        output = self._output / label
        if not output.exists():
            output.mkdir()
        output = output / (Path(self._image_path).stem
                           + '_{}_{}.jpg'.format(box[0], box[1]))
        cv2.imwrite(str(output), sub_image)
        with open(self._image_file, 'a') as file:
            file.write(str(output.relative_to(self._output))
                       + '&!&' + label + '\n')

    @abstractmethod
    def _update(self):
        """
        切割出所有的对象
        Returns
        -------

        """
        objects = self._image_info['objects'][:]
        shape = (self._image_info['shape']['height'],
                 self._image_info['shape']['width'])
        shape = np.asarray(shape).reshape([1, 2])
        labels = []
        boxes = []
        for ob in objects:
            labels.append(ob['name'])
            boxes.append(ob['box'])
        boxes = np.asarray(boxes)
        self._buf_data['boxes'] = boxes
        self._buf_data['shape'] = shape
        self._buf_data['labels'] = labels
        self._crop_image()
